package com.seo.proxy;

import java.io.File;
import java.io.IOException;
import java.net.InetSocketAddress;
import java.nio.ByteBuffer;
import java.nio.channels.SelectionKey;
import java.nio.channels.SocketChannel;

import org.apache.log4j.Logger;

import com.seo.async.AsyncSocket;
import com.seo.async.AsyncSocketListener;
import com.seo.async.HasClientContext;
import com.seo.async.Timer;
import com.seo.proxy.ban.BanManager;
import com.seo.proxy.context.ClientContext;
import com.seo.proxy.context.RequestContext;
import com.seo.proxy.context.ResponseContext;
import com.seo.proxy.http.Header;
import com.seo.proxy.http.HttpRequest;
import com.seo.proxy.http.HttpRequestParser;
import com.seo.proxy.replacer.UrlReplacer;

public class ProxyListener implements AsyncSocketListener {
	
	private final Logger LOGGER = Logger.getLogger(getClass());

	private static final int READ_WRITE_BUFFER_SIZE = 0x1000;
	
	private AsyncSocket asyncSocket;

	private ByteBuffer receiveBuffer = ByteBuffer.allocateDirect(0x1000);

	private HttpRequestParser requestParser = new HttpRequestParser();
	
	private RequestLogger requestLogger;
	
	// private Blacklist blacklist = new Blacklist();
	
	private BanManager banManager = new BanManager();
	
	private UrlReplacer urlReplacer = new UrlReplacer();
	
	public ProxyListener(File log) throws IOException {
		requestLogger = new RequestLogger(log);
	}

	@Override
	public void init(AsyncSocket asyncSocket) {
		this.asyncSocket = asyncSocket;
		
		asyncSocket.addTimer(10000, new Timer() {
			@Override
			public void onTimer() {
				banManager.unban();
			}
		});
	}

	@Override
	public Object onAccept(SocketChannel socketChannel) throws IOException {
		String ip = socketChannel.socket().getInetAddress().getHostAddress();
		if (banManager.isIpBanned(ip)) {
			asyncSocket.close(socketChannel);
			return null;
		}
		
		asyncSocket.subscribe(socketChannel, SelectionKey.OP_READ);
		ClientContext clientContext = new ClientContext(socketChannel, READ_WRITE_BUFFER_SIZE);
		clientContext.getRequestContext().setIp(ip);
		// LOGGER.info(ip + " connected");
		return clientContext.getRequestContext();
	}

	@Override
	public void onRead(SocketChannel socketChannel, Object context) throws IOException {
		if (context instanceof RequestContext) {
			readClientRequest(socketChannel, (RequestContext) context);
		} else {
			readProxyResponce(socketChannel, (ResponseContext) context);
		}
	}

	private void readClientRequest(SocketChannel socketChannel, RequestContext requestContext) throws IOException {
		receiveBuffer.clear();
		int size = socketChannel.read(receiveBuffer);
		if (size > 0) {
			Buffer requestBuffer = requestContext.getRequestBuffer();
			if (!requestBuffer.append(receiveBuffer)) {
				destroyContext(requestContext.getClientContext());
				return;
			}
			if (requestParser.parseRequest(requestContext)) {
				asyncSocket.unsubscribe(socketChannel);
				processClientRequest(requestContext);
			}
		} else if (size < 0) {
			destroyContext(requestContext.getClientContext());
		}
	}

	private void processClientRequest(RequestContext requestContext) throws IOException {
		/*String url = requestContext.getHttpRequest().getUrl().toString();
		if (blacklist.isBlacklisted(url)) {
			LOGGER.info("Blacklisted: " + url);
			destroyContext(requestContext.getClientContext());
			return;
		}*/
		
		int count = urlReplacer.replaceURLs(requestContext.getHttpRequest());
		banManager.request(requestContext.getHttpRequest().getUrl(), requestContext.getIp(), count > 0);
		
		LOGGER.info(requestContext.getSocketChannel().socket().getInetAddress().getHostAddress() + " " + requestContext.getHttpRequest().getRequestLine());
		requestLogger.logRequest(requestContext.getHttpRequest());
		
		HttpRequest httpRequest = requestContext.getHttpRequest();
		int port = httpRequest.getUrl().getPort();
		if (port == -1) {
			port = httpRequest.getUrl().getDefaultPort();
		}
		ResponseContext responseContext = new ResponseContext(requestContext.getClientContext());
		requestContext.getClientContext().setResponseContext(responseContext);
		SocketChannel socketChannel = asyncSocket.connect(new InetSocketAddress(httpRequest.getUrl().getHost(), port), responseContext);
		asyncSocket.subscribe(socketChannel, SelectionKey.OP_READ | SelectionKey.OP_WRITE | SelectionKey.OP_CONNECT);
	}

	private void readProxyResponce(SocketChannel socketChannel, ResponseContext responseContext) throws IOException {
		ClientContext clientContext = responseContext.getClientContext();
		ReadWriteBuffer readWriteBuffer = clientContext.getReadWriteBuffer();
		
		int size = readWriteBuffer.readToBufferFromChannel(socketChannel);
		if (size < 0) {
			asyncSocket.unsubscribe(socketChannel);
			responseContext.setProxyRequestDone(true);
			asyncSocket.subscribe(clientContext.getRequestContext().getSocketChannel(), SelectionKey.OP_WRITE);
			return;
		}
		if (readWriteBuffer.isFull()) {
			asyncSocket.unsubscribe(socketChannel);
		}
		if (readWriteBuffer.isDataAvailable()) {
			asyncSocket.subscribe(clientContext.getRequestContext().getSocketChannel(), SelectionKey.OP_WRITE);
		}
	}

	@Override
	public void onWrite(SocketChannel socketChannel, Object context) throws IOException {
		if (context instanceof RequestContext) {
			sendClientResponse(socketChannel, (RequestContext) context);
		} else {
			sendProxyRequest(socketChannel, (ResponseContext) context);
		}
	}

	private void sendProxyRequest(SocketChannel socketChannel, ResponseContext responseContext) throws IOException {
		if (responseContext.getProxyRequestBuffer() == null) {
			responseContext.setProxyRequestBuffer(createProxyRequestBuffer(responseContext));
		}
		socketChannel.write(responseContext.getProxyRequestBuffer());
		if (responseContext.getProxyRequestBuffer().remaining() == 0) {
			responseContext.setProxyRequestSent(true);
			asyncSocket.subscribe(socketChannel, SelectionKey.OP_READ);
		}
	}
	
	private ByteBuffer createProxyRequestBuffer(ResponseContext context) throws IOException {
		RequestContext requestContext = context.getClientContext().getRequestContext();
		HttpRequest httpRequest = requestContext.getHttpRequest();
		
		StringBuilder request = new StringBuilder();
		request.append(httpRequest.getMethod()).append(" ");
		request.append(httpRequest.getUrl().getFile()).append(" ");
		request.append(httpRequest.getProtocol()).append("\r\n");
		for (Header header : httpRequest.getHeaders()) {
			if (header.getName().equalsIgnoreCase("Connection")) {
				continue;
			}
			request.append(header.getName()).append(": ").append(header.getValue()).append("\r\n");
		}
		request.append("Connection: close\r\n\r\n");
		
		int size = request.length();
		if (httpRequest.getBody() != null) {
			size += httpRequest.getBody().length();
		}
		
		ByteBuffer buffer = ByteBuffer.allocate(size);
		buffer.put(request.toString().getBytes("iso-8859-1"));
		if (httpRequest.getBody() != null) {
			buffer.put(httpRequest.getBody().getBytes("iso-8859-1"));
		}
		buffer.flip();
		
		return buffer;
	}

	private void sendClientResponse(SocketChannel socketChannel, RequestContext requestContext) throws IOException {
		ClientContext clientContext = requestContext.getClientContext();
		ReadWriteBuffer readWriteBuffer = clientContext.getReadWriteBuffer();
		readWriteBuffer.writeFromBufferToChannel(socketChannel);
		if (readWriteBuffer.isFull()) {
			readWriteBuffer.compact();
			asyncSocket.subscribe(requestContext.getClientContext().getResponseContext().getSocketChannel(), SelectionKey.OP_READ);
		}
		if (!readWriteBuffer.isDataAvailable()) {
			if (clientContext.getResponseContext().isProxyRequestDone()) {
				if (!requestContext.getHttpRequest().isKeepConnection()) {
					destroyContext(clientContext);
				} else {
					clientContext.reset();
					if (!requestParser.parseRequest(requestContext)) {
						asyncSocket.subscribe(socketChannel, SelectionKey.OP_READ);
					} else {
						processClientRequest(requestContext);
					}
				}
			} else {
				asyncSocket.unsubscribe(socketChannel);
			}
		}
	}

	@Override
	public void onReadTimeout(SocketChannel socketChannel, Object context) {
		if (context instanceof RequestContext) {
			RequestContext requestContext = (RequestContext) context;
			LOGGER.info(requestContext.getIp() + " read timeout");
		} else {
			ResponseContext responseContext = (ResponseContext) context;
			RequestContext requestContext = responseContext.getClientContext().getRequestContext();
			LOGGER.info(responseContext.getClientContext().getRequestContext().getIp() + " proxy read timeout from " + requestContext.getHttpRequest().getUrl());
		}
		destroyContext(((HasClientContext) context).getClientContext());
	}

	@Override
	public void onWriteTimeout(SocketChannel socketChannel, Object context) {
		if (context instanceof RequestContext) {
			RequestContext requestContext = (RequestContext) context;
			LOGGER.info(requestContext.getIp() + " write timeout");
		} else {
			ResponseContext responseContext = (ResponseContext) context;
			LOGGER.info(responseContext.getClientContext().getRequestContext().getIp() + " proxy write timeout");
		}
		destroyContext(((HasClientContext) context).getClientContext());
	}

	private void destroyContext(ClientContext clientContext) {
		// LOGGER.info(clientContext.getRequestContext().getIp() + " closing connection");
		asyncSocket.close(clientContext.getRequestContext().getSocketChannel());
		if (clientContext.getResponseContext() != null) {
			asyncSocket.close(clientContext.getResponseContext().getSocketChannel());
		}
	}

	@Override
	public void onError(SocketChannel socketChannel, Object context) {
		if (context != null) {
			destroyContext(((HasClientContext) context).getClientContext());
		}
	}
}
